// 图片下载器
const { chromium } = require('playwright');
const fs = require('fs');
const path = require('path');
const readline = require('readline');

// 创建一个读取用户输入的接口
const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout
});

// 提示用户输入，返回Promise
function promptUser(question) {
  return new Promise((resolve) => {
    rl.question(question, (answer) => {
      resolve(answer);
    });
  });
}

// 创建存储目录
function createDirectory(directoryPath) {
  if (!fs.existsSync(directoryPath)) {
    fs.mkdirSync(directoryPath, { recursive: true });
    console.log(`创建目录: ${directoryPath}`);
  }
}

// 收集当前页面的图片链接
async function collectImagesFromCurrentPage(page) {
  try {
    return await page.evaluate(() => {
      const imageLinks = [];
      
      // 尝试多种图片选择器
      const selectors = [
        'img.content_img', // 内容图片
        '.content img',    // 内容区域中的图片
        '.picbox img',     // 图片框中的图片
        '.article-content img', // 文章内容中的图片
        '.detail-content img',  // 详情内容中的图片
        'div[id*="content"] img',  // ID包含content的div中的图片
        '.pic img', // 图片容器中的图片
        '.photo img', // 照片容器中的图片
        'img[src*=".jpg"]', // src中包含.jpg的图片
        'img[src*=".jpeg"]',
        'img[src*=".png"]',
        'img[src*=".webp"]', // 添加对webp格式的支持
        'img[src*=".gif"]',  // 添加对gif格式的支持
        'img[src*=".avif"]', // 添加对avif格式的支持
        'img[data-original]', // 懒加载图片
        'img[data-src]',      // 另一种懒加载图片
        'img[loading="lazy"]' // HTML5懒加载属性
      ];
      
      for (const selector of selectors) {
        try {
          const elements = document.querySelectorAll(selector);
          if (elements.length > 0) {
            for (const el of elements) {
              // 优先获取原始图片链接（懒加载）
              let imgSrc = el.getAttribute('data-original') || 
                           el.getAttribute('data-src') || 
                           el.getAttribute('src');
              
              if (imgSrc) {
                // 如果是相对链接，转换为绝对链接
                if (imgSrc.startsWith('/')) {
                  imgSrc = window.location.origin + imgSrc;
                } else if (!imgSrc.startsWith('http')) {
                  const base = window.location.href.substring(0, window.location.href.lastIndexOf('/') + 1);
                  imgSrc = base + imgSrc;
                }
                
                // 过滤掉广告图片、图标等小图片
                if (!imgSrc.includes('icon') && 
                    !imgSrc.includes('logo') && 
                    !imgSrc.includes('banner') &&
                    !imgSrc.includes('ad.') &&
                    !imgSrc.includes('favicon')) {
                  // 检查图片尺寸（如果有）
                  const width = el.getAttribute('width') || el.width;
                  const height = el.getAttribute('height') || el.height;
                  
                  // 如果图片足够大，或者没有尺寸信息，则添加到链接列表
                  if (!width || !height || (parseInt(width) > 200 && parseInt(height) > 200)) {
                    imageLinks.push(imgSrc);
                  }
                }
              }
            }
          }
        } catch (err) {
          console.error(`选择器 ${selector} 出错:`, err);
        }
      }
      
      return imageLinks;
    });
  } catch (error) {
    console.error('收集图片链接时出错:', error);
    return [];
  }
}

async function main() {
  try {
    // 从环境变量或用户输入获取模特名称
    let modelName = process.env.MODEL_NAME;
    let isHeadless = process.env.HEADLESS_MODE === 'y';
    let selectedIndices = process.env.SELECTED_INDICES;
    
    // 如果环境变量未设置，则询问用户
    if (!modelName) {
      modelName = await promptUser('请输入模特名称: ');
      if (!modelName.trim()) {
        console.log('模特名称不能为空');
        rl.close();
        return;
      }
    }
    
    // 检查链接文件是否存在
    const linksFileName = `${modelName}_links.json`;
    if (!fs.existsSync(linksFileName)) {
      console.log(`链接文件 ${linksFileName} 不存在，请先运行爬虫程序获取链接`);
      rl.close();
      return;
    }
    
    // 加载链接文件
    const links = JSON.parse(fs.readFileSync(linksFileName, 'utf8'));
    console.log(`找到了 ${links.length} 个图片集链接`);
    
    // 显示所有链接
    links.forEach((link, index) => {
      console.log(`[${index + 1}] ${link.title} (${link.url})`);
    });
    
    // 如果环境变量未设置选择的索引，则询问用户
    let indices = [];
    if (!selectedIndices) {
      // 询问用户选择要下载的图片集
      const answer = await promptUser('请输入要下载的图片集序号(用逗号分隔)，或输入"all"下载全部: ');
      if (answer.trim().toLowerCase() === 'all') {
        indices = links.map((_, index) => index);
      } else {
        // 转换输入的序号为索引数组
        indices = answer.split(',')
          .map(num => parseInt(num.trim()) - 1)
          .filter(index => !isNaN(index) && index >= 0 && index < links.length);
      }
    } else {
      // 解析环境变量中的选择索引
      if (selectedIndices.trim().toLowerCase() === 'all') {
        indices = links.map((_, index) => index);
      } else {
        indices = selectedIndices.split(',')
          .map(num => parseInt(num.trim()) - 1)
          .filter(index => !isNaN(index) && index >= 0 && index < links.length);
      }
    }
    
    if (indices.length === 0) {
      console.log('没有选择有效的图片集，任务退出');
      rl.close();
      return;
    }
    
    console.log(`将下载 ${indices.length} 个图片集`);
    
    // 如果环境变量未设置无头模式，则询问用户
    if (process.env.HEADLESS_MODE === undefined) {
      const headlessMode = await promptUser('是否使用无头模式(不显示浏览器)？(y/n，默认y): ');
      isHeadless = headlessMode.trim().toLowerCase() !== 'n';
    }
    
    console.log(`浏览器模式: ${isHeadless ? '无头模式' : '有头模式'}`);
    
    // 创建浏览器实例
    const browser = await chromium.launch({
      headless: isHeadless
    });
    
    // 创建存储目录
    const baseDir = path.join(__dirname, modelName);
    createDirectory(baseDir);
    
    // 创建Page对象
    const page = await browser.newPage();
    
    // 下载选择的图片集
    for (let i = 0; i < indices.length; i++) {
      const index = indices[i];
      const link = links[index];
      
      console.log(`\n[${i + 1}/${indices.length}] 正在处理图片集: ${link.title}`);
      
      // 确保URL包含主机名
      let url = link.url;
      if (!url.startsWith('http')) {
        url = 'http://3c.xy05.my' + url;
      }
      
      try {
        // 访问图片集页面
        console.log(`访问URL: ${url}`);
        await page.goto(url, { timeout: 60000 });
        await page.waitForLoadState('networkidle');
        
        // 获取图片集标题，使用URL中的标题作为备用
        let title = link.title;
        if (!title) {
          const titleElement = await page.$('.content h1');
          if (titleElement) {
            title = await titleElement.innerText();
          } else {
            // 从URL中提取一个标题
            const urlParts = url.split('/');
            title = urlParts[urlParts.length - 1].split('.')[0] || `图片集_${index}`;
          }
        }
        
        // 创建图片集目录
        const safeTitle = title.replace(/[\\/:*?"<>|]/g, '_');
        const collectionDir = path.join(baseDir, safeTitle);
        createDirectory(collectionDir);
        
        // 获取所有图片链接
        const imgLinks = await collectImagesFromCurrentPage(page);
        
        console.log(`找到了 ${imgLinks.length} 张图片`);
        
        // 下载所有图片
        for (let j = 0; j < imgLinks.length; j++) {
          try {
            let imgUrl = imgLinks[j];
            
            // 确保URL包含主机名
            if (!imgUrl.startsWith('http')) {
              imgUrl = 'http://3c.xy05.my' + imgUrl;
            }
            
            // 构建图片文件名
            const imgFileName = `${j + 1}.jpg`;
            const imgPath = path.join(collectionDir, imgFileName);
            
            // 如果文件已存在，跳过下载
            if (fs.existsSync(imgPath)) {
              console.log(`[${j + 1}/${imgLinks.length}] 图片已存在，跳过: ${imgFileName}`);
              continue;
            }
            
            console.log(`[${j + 1}/${imgLinks.length}] 下载图片: ${imgUrl}`);
            
            // 获取图片内容
            const imgResponse = await page.goto(imgUrl, { timeout: 30000 });
            const imgBuffer = await imgResponse.body();
            
            // 保存图片
            fs.writeFileSync(imgPath, imgBuffer);
            
            // 返回原页面
            await page.goBack();
            await page.waitForLoadState('networkidle');
            
          } catch (imgError) {
            console.error(`下载图片时发生错误: ${imgError.message}`);
          }
        }
        
        console.log(`图片集 "${title}" 下载完成`);
        
      } catch (pageError) {
        console.error(`处理图片集时发生错误: ${pageError.message}`);
      }
    }
    
    console.log('\n所有图片集下载完成！');
    
    // 关闭浏览器
    await browser.close();
    
    // 关闭readline接口
    rl.close();
    
  } catch (error) {
    console.error('发生错误:', error);
    rl.close();
  }
}

// 执行主函数
main().catch(console.error); 